from PIL import Image
import matplotlib.pyplot as plt
import torch
from torchvision import transforms
from transformers import AutoModelForImageSegmentation
from tqdm import tqdm
import pathlib
import multiprocessing
import sys

class RemoveBG:
  def __init__(self, local_files_only=False):
    self.model = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True, local_files_only=local_files_only)
    torch.set_float32_matmul_precision(['high', 'highest'][0])
    self.model.to('cpu')
    self.model.eval()

    # Data settings
    image_size = (1024, 1024)
    self.transform_image = transforms.Compose([
        transforms.Resize(image_size),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

  def __call__(self, input_image_path):
    image = Image.open(input_image_path)
    input_images = self.transform_image(image).unsqueeze(0).to('cpu')

    # Prediction
    with torch.no_grad():
        preds = self.model(input_images)[-1].sigmoid().cpu()
    pred = preds[0].squeeze()
    pred_pil = transforms.ToPILImage()(pred)
    mask = pred_pil.resize(image.size)
    image.putalpha(mask)
    return image
  
def remove_many_bg(input_image_paths, output_image_directory, local_files_only=True):
    remove_bg = RemoveBG(local_files_only=local_files_only)
    if not output_image_directory:
        use_output_image_directory = True
    else:
        use_output_image_directory = False
    for input_image_path in tqdm(input_image_paths):
        image = remove_bg(input_image_path)
        # If output_image_directory is not specified, then use the same directory as input_image_path
        if use_output_image_directory:
            output_image_directory = str(pathlib.Path(input_image_path).parent)
        # Create output directory if not exists
        pathlib.Path(output_image_directory).mkdir(parents=True, exist_ok=True)
        output_image_fullpath = pathlib.Path(output_image_directory) / pathlib.Path(input_image_path).name.replace(".jpg", ".png").replace(".jpeg", ".png")
        image.save(output_image_fullpath)


if __name__ == '__main__':
  # Check for -s command line argument (skip user input)
  if '-s' in sys.argv:
    input_image_path = ''
  else:
    input_image_path = input('Please input image path (default read from "input_images.txt"): ').strip().replace('"', '').replace("'", '')
  if input_image_path:
    #If input_image_path is a directory, then list all jpg files in the directory and remove background
    if pathlib.Path(input_image_path).is_dir():
      input_image_paths = [str(pathlib.Path(input_image_path) / pathlib.Path(file)) for file in pathlib.Path(input_image_path).glob('*.jpg')]
    else:
      input_image_paths = [input_image_path]
  else:
    try:
      # First try to read from utf8
      with open('input_images.txt', 'r', encoding='utf8') as f:
        input_image_paths = [line.strip()[1:-1] if line.strip().startswith('"') or line.strip().startswith("'") else line.strip() for line in f.readlines() if line.strip()]
    except:
      try:
        # Then try to read from gbk
        with open('input_images.txt', 'r') as f:
          input_image_paths = [line.strip()[1:-1] if line.strip().startswith('"') or line.strip().startswith("'") else line.strip() for line in f.readlines() if line.strip()]
      except:
        print('input_images.txt not found or unknown encoding!')
        input_image_paths = []

  if '-s' in sys.argv:
    output_image_directory = ''
  else:
    output_image_directory = input('Please input output image directory (default directory same as input_images): ').strip()
  
  if len(input_image_paths) > 0:
    if '-s' in sys.argv:
        local_files_only = True
    elif input('Do you want to fetch latest model? (y/n) ').strip().lower() == 'y':
      print('Fetching latest model...')
      local_files_only = False
    else:
      local_files_only = True
    remove_many_bg(input_image_paths, output_image_directory, local_files_only=local_files_only)
